'''
Author: SlytherinGe
LastEditTime: 2021-07-02 16:24:08
'''
import numpy as np
import cv2

if __name__ == '__main__':
    from mmdet.datasets.builder import PIPELINES
else:
    from ..builder import PIPELINES

SMALL_NUM = 1e-12

@PIPELINES.register_module()
class PAOIGenerate(object):

    def __init__(self,
                 sigma,
                 pos_thr=1.0,
                 sigma_ratio=1.0):
        super().__init__()

        self.sigma = sigma
        self.pos_thr = pos_thr
        self.sigma_ratio = sigma_ratio

    def __call__(self, results):
        # get bboxes data from results['gt_bboxes'], which is a numpy array
        # get img shape from results['img_shape'], which is a tuple(h, w, c)
        gt_bboxes = results['gt_bboxes'].astype(np.intp)
        img_shape = results['img_shape']
        if self.sigma is None:
            sigmax = np.ones((gt_bboxes.shape[0], 2), dtype=np.int)
            sigmay = np.ones((gt_bboxes.shape[0], 2), dtype=np.int)
            sigmax[:,0] = (gt_bboxes[:,2]-gt_bboxes[:,0])*self.sigma_ratio
            sigmay[:,0] = (gt_bboxes[:,3]-gt_bboxes[:,1])*self.sigma_ratio
            sigmax = np.max(sigmax, axis=-1)
            sigmay = np.max(sigmay, axis=-1)

        background = np.zeros((gt_bboxes.shape[0], img_shape[0], img_shape[1]))
        for i in range(gt_bboxes.shape[0]):
            background[i,gt_bboxes[i,1]:gt_bboxes[i,3],gt_bboxes[i,0]:gt_bboxes[i,2]] = 1.0
            if self.sigma is not None:
                background[i] = cv2.GaussianBlur(background[i],(0,0),self.sigma)
            else:
                background[i] = cv2.GaussianBlur(background[i],(0,0),sigmaX=int(sigmax[i]),sigmaY=int(sigmay[i]))
            background[i] /= (np.max(background[i]) + SMALL_NUM)
        heat_map = np.sum(background,axis=0)
        heat_map[heat_map>self.pos_thr] = 1.0

        results['gt_masks'] = heat_map
        return results        
